/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.schemas.transforms;



import com.google.auto.value.AutoValue;

import com.bff.gaia.unified.sdk.annotations.Experimental;

import com.bff.gaia.unified.sdk.schemas.FieldAccessDescriptor;

import com.bff.gaia.unified.sdk.schemas.Schema;

import com.bff.gaia.unified.sdk.schemas.utils.SchemaZipFold;

import com.bff.gaia.unified.sdk.transforms.DoFn;

import com.bff.gaia.unified.sdk.transforms.PTransform;

import com.bff.gaia.unified.sdk.transforms.ParDo;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.sdk.values.Row;

import com.bff.gaia.unified.vendor.guava.com.google.common.base.Joiner;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableList;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Maps;



import java.io.Serializable;

import java.math.BigDecimal;

import java.util.*;

import java.util.stream.Collectors;



/** Set of utilities for casting rows between schemas. */

@Experimental(Experimental.Kind.SCHEMAS)

@AutoValue

public abstract class Cast<T> extends PTransform<PCollection<T>, PCollection<Row>> {



  public abstract Schema outputSchema();



  public abstract Validator validator();



  public static <T> Cast<T> of(Schema outputSchema, Validator validator) {

    return new AutoValue_Cast<>(outputSchema, validator);

  }



  public static <T> Cast<T> widening(Schema outputSchema) {

    return new AutoValue_Cast<>(outputSchema, Widening.of());

  }



  public static <T> Cast<T> narrowing(Schema outputSchema) {

    return new AutoValue_Cast<>(outputSchema, Narrowing.of());

  }



  /** Describes compatibility errors during casting. */

  @AutoValue

  public abstract static class CompatibilityError implements Serializable {



    public abstract List<String> path();



    public abstract String message();



    public static CompatibilityError create(List<String> path, String message) {

      return new AutoValue_Cast_CompatibilityError(path, message);

    }

  }



  /** Interface for statically validating casts. */

  public interface Validator extends Serializable {

    List<CompatibilityError> apply(Schema input, Schema output);

  }



  /**

   * Widening changes to type that can represent any possible value of the original type.

   *

   * <p>Standard widening conversions:

   *

   * <ul>

   *   <li>BYTE to INT16, INT32, INT64, FLOAT, DOUBLE, DECIMAL

   *   <li>INT16 to INT32, INT64, FLOAT, DOUBLE, DECIMAL

   *   <li>INT32 to INT64, FLOAT, DOUBLE, DECIMAL

   *   <li>INT64 to FLOAT, DOUBLE, DECIMAL

   *   <li>FLOAT to DOUBLE, DECIMAL

   *   <li>DOUBLE to DECIMAL

   * </ul>

   *

   * <p>Row widening:

   *

   * <ul>

   *   <li>wider schema to schema with a subset of fields

   *   <li>non-nullable fields to nullable fields

   * </ul>

   *

   * <p>Widening doesn't lose information about the overall magnitude in following cases:

   *

   * <ul>

   *   <li>integral type to another integral type

   *   <li>BYTE or INT16 to FLOAT, DOUBLE or DECIMAL

   *   <li>INT32 to DOUBLE

   * </ul>

   *

   * <p>Other conversions to may cause loss of precision.

   */

  public static class Widening implements Validator {

    private final Fold fold = new Fold();



    public static Widening of() {

      return new Widening();

    }



    @Override

    public String toString() {

      return "Cast.Widening";

    }



    @Override

    public List<CompatibilityError> apply(final Schema input, final Schema output) {

      return fold.apply(input, output);

    }



    private static class Fold extends SchemaZipFold<List<CompatibilityError>> {



      @Override

      public List<CompatibilityError> accumulate(

          List<CompatibilityError> left, List<CompatibilityError> right) {

        return ImmutableList.<CompatibilityError>builder().addAll(left).addAll(right).build();

      }



      @Override

      public List<CompatibilityError> accept(

		  Context context, Optional<Schema.Field> left, Optional<Schema.Field> right) {

        if (!left.isPresent() && !right.isPresent()) {

          return Collections.emptyList();

        } else if (left.isPresent() && !right.isPresent()) {

          return Collections.emptyList();

        } else if (!left.isPresent() && right.isPresent()) {

          return Collections.singletonList(

              CompatibilityError.create(context.path(), "Field is missing in output schema"));

        } else {

          if (left.get().getType().getNullable() && !right.get().getType().getNullable()) {

            return Collections.singletonList(

                CompatibilityError.create(

                    context.path(), "Can't cast nullable field to non-nullable field"));

          }

        }



        return Collections.emptyList();

      }



      @Override

      public List<CompatibilityError> accept(Context context, Schema.FieldType input, Schema.FieldType output) {

        Schema.TypeName inputType = input.getTypeName();

        Schema.TypeName outputType = output.getTypeName();



        boolean supertype = outputType.isSupertypeOf(inputType);



        if (isIntegral(inputType) && isDecimal(outputType)) {

          return Collections.emptyList();

        } else if (!supertype) {

          return Collections.singletonList(

              CompatibilityError.create(

                  context.path(), "Can't cast '" + inputType + "' to '" + outputType + "'"));

        }



        return Collections.emptyList();

      }

    }

  }



  /**

   * Narrowing changes type without guarantee to preserve data.

   *

   * <p>Standard narrowing conversions:

   *

   * <ul>

   *   <li>any conversions of {@link Widening}

   *   <li>conversions the opposite to {@link Widening}

   * </ul>

   *

   * <p>Row narrowing

   *

   * <ul>

   *   <li>wider schema to schema with a subset of fields

   *   <li>non-nullable fields to nullable fields

   *   <li>nullable fields to non-nullable fields

   * </ul>

   */

  public static class Narrowing implements Validator {

    private final Fold fold = new Fold();



    public static Narrowing of() {

      return new Narrowing();

    }



    @Override

    public String toString() {

      return "Cast.Narrowing";

    }



    @Override

    public List<CompatibilityError> apply(final Schema input, final Schema output) {

      return fold.apply(input, output);

    }



    private static class Fold extends SchemaZipFold<List<CompatibilityError>> {



      @Override

      public List<CompatibilityError> accumulate(

          List<CompatibilityError> left, List<CompatibilityError> right) {

        return ImmutableList.<CompatibilityError>builder().addAll(left).addAll(right).build();

      }



      @Override

      public List<CompatibilityError> accept(

		  Context context, Optional<Schema.Field> left, Optional<Schema.Field> right) {



        if (!left.isPresent() && right.isPresent()) {

          return Collections.singletonList(

              CompatibilityError.create(context.path(), "Field is missing in output schema"));

        }



        return Collections.emptyList();

      }



      @Override

      public List<CompatibilityError> accept(Context context, Schema.FieldType input, Schema.FieldType output) {

        Schema.TypeName inputType = input.getTypeName();

        Schema.TypeName outputType = output.getTypeName();



        boolean supertype = outputType.isSupertypeOf(inputType);

        boolean subtype = outputType.isSubtypeOf(inputType);



        if (isDecimal(inputType) && isIntegral(outputType)) {

          return Collections.emptyList();

        } else if (!supertype && !subtype) {

          return Collections.singletonList(

              CompatibilityError.create(

                  context.path(), "Can't cast '" + inputType + "' to '" + outputType + "'"));

        }



        return Collections.emptyList();

      }

    }

  }



  /** Checks if type is integral. */

  public static boolean isIntegral(Schema.TypeName type) {

    return type == Schema.TypeName.BYTE

        || type == Schema.TypeName.INT16

        || type == Schema.TypeName.INT32

        || type == Schema.TypeName.INT64;

  }



  /** Checks if type is decimal. */

  public static boolean isDecimal(Schema.TypeName type) {

    return type == Schema.TypeName.FLOAT || type == Schema.TypeName.DOUBLE || type == Schema.TypeName.DECIMAL;

  }



  public void verifyCompatibility(Schema inputSchema) {

    List<CompatibilityError> errors = validator().apply(inputSchema, outputSchema());



    if (!errors.isEmpty()) {

      String reason =

          errors.stream()

              .map(x -> Joiner.on('.').join(x.path()) + ": " + x.message())

              .collect(Collectors.joining("\n\t"));



      throw new IllegalArgumentException(

          "Cast isn't compatible using " + validator() + ":\n\t" + reason);

    }

  }



  @Override

  public PCollection<Row> expand(PCollection<T> input) {

    Schema inputSchema = input.getSchema();



    verifyCompatibility(inputSchema);



    return input

        .apply(

            ParDo.of(

                new DoFn<T, Row>() {

                  // TODO: This should be the same as resolved so that Unified knows which fields

                  // are being accessed. Currently Unified only supports wildcard descriptors.

                  // Once BEAM-4457 is fixed, fix this.

                  @FieldAccess("filterFields")

                  final FieldAccessDescriptor fieldAccessDescriptor =

                      FieldAccessDescriptor.withAllFields();



                  @ProcessElement

                  public void process(

					  @FieldAccess("filterFields") @Element Row input, OutputReceiver<Row> r) {

                    Row output = castRow(input, inputSchema, outputSchema());

                    r.output(output);

                  }

                }))

        .setRowSchema(outputSchema());

  }



  public static Row castRow(Row input, Schema inputSchema, Schema outputSchema) {

    if (input == null) {

      return null;

    }



    Row.Builder output = Row.withSchema(outputSchema);

    for (int i = 0; i < outputSchema.getFieldCount(); i++) {

      Schema.Field outputField = outputSchema.getField(i);



      int fromFieldIdx = inputSchema.indexOf(outputField.getName());

      Schema.Field inputField = inputSchema.getField(fromFieldIdx);



      Object inputValue = input.getValue(fromFieldIdx);

      Object outputValue = castValue(inputValue, inputField.getType(), outputField.getType());



      output.addValue(outputValue);

    }



    return output.build();

  }



  public static Number castNumber(Number value, Schema.TypeName input, Schema.TypeName output) {

    if (!input.isNumericType()) {

      throw new RuntimeException("Can't cast non-numeric types: " + input);

    }



    if (!output.isNumericType()) {

      throw new RuntimeException("Can't cast numbers to non-numeric type: " + output);

    }



    if (value == null) {

      return null;

    }



    if (input == output) {

      return value;

    }



    switch (output) {

      case BYTE:

        return value.byteValue();



      case INT16:

        return value.shortValue();



      case INT32:

        return value.intValue();



      case INT64:

        return value.longValue();



      case FLOAT:

        return value.floatValue();



      case DOUBLE:

        return value.doubleValue();



      case DECIMAL:

        switch (input) {

          case BYTE:

          case INT16:

          case INT32:

            return new BigDecimal(value.intValue());



          case INT64:

            return new BigDecimal(value.longValue());



          case FLOAT:

          case DOUBLE:

            return new BigDecimal(value.doubleValue());



          default:

            throw new AssertionError("Unexpected numeric type: " + output);

        }



      default:

        throw new AssertionError("Unexpected numeric type: " + output);

    }

  }



  @SuppressWarnings("unchecked")

  public static Object castValue(Object inputValue, Schema.FieldType input, Schema.FieldType output) {



    Schema.TypeName inputType = input.getTypeName();

    Schema.TypeName outputType = output.getTypeName();



    if (inputValue == null) {

      return null;

    }



    switch (inputType) {

      case ROW:

        return castRow((Row) inputValue, input.getRowSchema(), output.getRowSchema());



      case ARRAY:

        List<Object> inputValues = (List<Object>) inputValue;

        List<Object> outputValues = new ArrayList<>(inputValues.size());



        for (Object elem : inputValues) {

          outputValues.add(

              castValue(elem, input.getCollectionElementType(), output.getCollectionElementType()));

        }



        return outputValues;



      case MAP:

        Map<Object, Object> inputMap = (Map<Object, Object>) inputValue;

        Map<Object, Object> outputMap = Maps.newHashMapWithExpectedSize(inputMap.size());



        for (Map.Entry<Object, Object> entry : inputMap.entrySet()) {

          Object outputKey =

              castValue(entry.getKey(), input.getMapKeyType(), output.getMapKeyType());

          Object outputValue =

              castValue(entry.getValue(), input.getMapValueType(), output.getMapValueType());



          outputMap.put(outputKey, outputValue);

        }



        return outputMap;



      default:

        if (inputType.equals(outputType)) {

          return inputValue;

        }



        if (inputType.isNumericType()) {

          return castNumber((Number) inputValue, inputType, outputType);

        } else {

          throw new IllegalArgumentException("input should be array, map, numeric or row");

        }

    }

  }

}